from typing import Optional, Union, Sequence

import moderngl
import pyrr

from . import geometry
from . import material
from .attribute import Attribute, Constant
from .camera import Camera
from .light import LightManager


class Object3D:
    def __init__(self, ctx: moderngl.Context, geo: geometry.Geometry, mat: material.Material):
        self.geometry = geo
        self.material = mat
        self.ctx = ctx
        self._position: pyrr.Vector4 = pyrr.Vector4((0, 0, 0, 1), dtype='f4')
        self._matrix: Optional[pyrr.Matrix44] = None
        self._trans_matrix: pyrr.Matrix44 = pyrr.Matrix44.identity()
        self._rotate_matrix: pyrr.Matrix44 = pyrr.Matrix44.identity()
        self._program = ctx.program(**mat.shaders)
        attrs = []
        for attr_const in mat.attributes.keys():
            mat_attr = mat.attributes[attr_const]
            geo_attr = geo.attributes[attr_const]
            assert Attribute.is_same_type(mat_attr, geo_attr)
            buffer = ctx.buffer(getattr(geo, geo_attr.name))
            attrs.append((buffer, *mat_attr))
        if Constant.INDEX in geo.attributes:
            index_attr = geo.attributes[Constant.INDEX]
            ibo = ctx.buffer(getattr(geo, index_attr.name))
            self._vao = ctx.vertex_array(self._program, attrs, ibo)
        else:
            self._vao = ctx.vertex_array(self._program, attrs)

    @property
    def matrix(self) -> pyrr.Matrix44:
        if self._matrix is None:
            self._matrix = self._trans_matrix * self._rotate_matrix
        return self._matrix

    def rotate_y(self, angle: float):
        rotate_matrix = pyrr.Matrix44.from_y_rotation(angle)
        self._rotate_matrix = rotate_matrix * self._rotate_matrix
        self._matrix = None

    def translate(
            self, x: Union[float, Sequence[float]] = 0,
            y: float = 0, z: float = 0
    ):
        if isinstance(x, Sequence):
            trans = pyrr.Vector3(x)
        else:
            trans = pyrr.Vector3((x, y, z))
        trans_matrix = pyrr.Matrix44.from_translation(trans)
        self._position = trans_matrix * self._position
        self._trans_matrix = trans_matrix * self._trans_matrix
        self._matrix = None

    def set_position(
            self, x: Union[float, Sequence[float]] = 0,
            y: float = None, z: float = None
    ):
        if isinstance(x, Sequence):
            new_pos = pyrr.Vector4((*x, 1), dtype='f4')
        else:
            if x is None:
                x = self._position.x
            if y is None:
                y = self._position.y
            if z is None:
                z = self._position.z
            new_pos = pyrr.Vector4((x, y, z, 1), dtype='f4')
        self._trans_matrix = pyrr.Matrix44.from_translation(new_pos.xyz)
        self._matrix = None

    def get_position(self) -> pyrr.Vector3:
        return self._position.xyz

    def render(self, camera: Camera, light_manager: LightManager):
        vao = self._vao
        program = self._program
        geo = self.geometry
        mat = self.material
        for uniform_key, uniform_value in mat.uniforms:
            program[uniform_key] = uniform_value
        for uniform_key, uniform_value in light_manager.uniforms:
            program[uniform_key] = uniform_value
        program['cameraMatrix'] = camera.matrix.flatten()
        program['modelMatrix'] = self.matrix.flatten()
        vao.render(mode=geo.mode)
